import numpy as np
import pandas as pd
import math
import matplotlib
matplotlib.use('TkAgg')  # 或者使用其他后端 'Agg', 'Qt5Agg' 等
import matplotlib.pyplot as plt
import os
# 地球半径，单位为米
R = 6378137.0


# 将度数转换为弧度
def deg2rad(deg):
    return deg * (math.pi / 180)


# 批量转换经纬度数据为x, y直角坐标
def PVT2ENU(locations, base_loc):
    # 将纬度和经度差转换为弧度
    delta_lat = deg2rad(locations[:, 0] - base_loc[0])  # 纬度差
    delta_lon = deg2rad(locations[:, 1] - base_loc[1])  # 经度差

    # 基准点的纬度转换为弧度
    lat_base_rad = deg2rad(base_loc[0])

    # 计算Y（北向坐标）
    y = R * delta_lat

    # 计算X（东向坐标），需要乘以纬度的余弦
    x = R * delta_lon * np.cos(lat_base_rad)

    # 返回 (x, y) 坐标
    return np.stack((x, y), axis=1)

def sensor_interp(tmsp, sdata):
    sync_cols = []
    for col_idx in range(sdata.shape[1]):
       sync_sdata = np.interp(tmsp, sdata[:, 0], sdata[:, col_idx])
       sync_cols.append(sync_sdata)
    sync_sdata = np.stack(sync_cols, axis=1)
    return sync_sdata

def get_data(head):
    loca_path = os.path.join(head,'Location.csv')
    acce_path = os.path.join(head,'Accelerometer.csv')
    gyro_path = os.path.join(head,'Gyroscope.csv')
    orie_path = os.path.join(head,'Orientation.csv')

    location_data = pd.read_csv(loca_path)
    acce_data = pd.read_csv(acce_path)[['time','x','y','z']].to_numpy()
    gyro_data = pd.read_csv(gyro_path)[['x','y','z']].to_numpy()
    orie_data = pd.read_csv(orie_path)[['qw','qx','qy','qz']].to_numpy()
    '''
    test data 
    '''
    '''imu_path = os.path.join(head, 'imu1.csv')
    imu_data = pd.read_csv(imu_path)

    #  tmsp,accx,accy,accz,gyrox,gyroy,gyroz,rotx,roty,rotz,rotw
    acce_data = imu_data[['tmsp', 'accx', 'accy', 'accz']].to_numpy()
    gyro_data = imu_data[['gyrox', 'gyroy', 'gyroz']].to_numpy()
    orie_data = imu_data[['rotw', 'rotx', 'roty', 'rotz']].to_numpy()'''

    imu_all = np.concatenate((acce_data, gyro_data), axis=1)
    time_data = imu_all[:,0]
    gps_data = []
    base_loc = []
    time_stmp = []
    for ti, la, lo in zip(location_data['time'], location_data['latitude'], location_data['longitude']):
        gps_data.append([ti, la, lo])

    for i in range(len(time_data)):
        time_stmp.append(time_data[i])
    gps_data = np.array(gps_data)
    time_stmp = np.array(time_stmp)
    base_loc = gps_data[0][1:]

    # interpolation
    interp_data = sensor_interp(time_stmp, gps_data)
    interp_data = interp_data[:, 1:]
    xy_point = PVT2ENU(interp_data, base_loc)
    print(imu_all.shape)
    print(xy_point.shape)
    print(orie_data.shape)
    return imu_all, xy_point, orie_data
    # 计算路径
#path = '/home/hulab/competition/ronin-master_new/source/preprocessing/Data/my_raw_data/_-2024-10-04_05-50-31/'
#get_data(path)
